rm(list = ls())
library(reshape2)
library(ggplot2)
library(magrittr)
library(stringr)
devtools::load_all()

# ---------函数--------
# GothVP is defined as 后一期值函数的导数
GothVP <- function(a, tranval,permval,Dep,Beta,R,W,M,cons){

  # EUP 是效用函数导数的期望
  EUP <- matlab::zeros(matlab::size(a))

  # 这里的循环主要是为了计算期望，需要遍历各个概率
  for (i in 1:length(tranval)){
    for (j in 1:length(permval)){
      CapPsi <- permval[j]
      CapXi <- tranval[i]
      k <- a*Dep/(G*CapPsi)
      EUP <- EUP+Dep*Beta*R(k,CurlyEpsilon)*
        UP(G*CapPsi* signal::interp1(M,cons,R(k,CurlyEpsilon)*k + W(k,CurlyEpsilon)*CapXi,extrap = T),Rho)*
        tranprob[i]*permprob[j]
    }
  }
  return(EUP)
}

NP <- function(cc, Rho) cc^(-1/Rho) # u'的反函数
UP <- function(cc,Rho) cc^(-Rho) # u'

# --------求解----------
# 参数设置
Beta <- 0.96   # time preference
Rho <- 2       # coefficient of relative risk aversion
n <- 20        # number of grid points
G<-1.01      # permanent income growth rate
Dep<-0.90    # depreciation rate
CurlyEpsilon<-0.36  # Cobb-Douglas production function parameter
permval<-c(0.90,1.00,1.10)   # permanent shock values
permprob<-c(.25,0.50,0.25)  # permanent shock probabilities
tranval<-1.00             # transitory  shock values
tranprob<-1.00            # transitoty shock probabilities
R <-function(k,CurlyEpsilon) 1+CurlyEpsilon*k^(CurlyEpsilon-1)   # set up interest rate function
W <- function(k,CurlyEpsilon) (1-CurlyEpsilon)*(k^CurlyEpsilon)   # set up wage rate function
kSS <- ((((G^Rho)/(Beta*Dep))-1)/CurlyEpsilon)^(1/(CurlyEpsilon-1))
aSS <- kSS*G/Dep
AlphaVec <- exp(seq(0,log(3*aSS),length.out = 20))-1       # set up grid points
PeriodsToAdd=99 # 迭代期数


# 初始化
cons <- seq(0,n-1) %>% matrix(ncol = 1)
M <- seq(0,n-1) %>% matrix(ncol = 1)

# 求解
for (i in 1:PeriodsToAdd){
  # calculate ct from each grid point in AlphaVec
  ChiVec <- NP(GothVP(AlphaVec,tranval,permval,Dep,Beta,R,W,M[,ncol(M)],cons[,ncol(cons)]),Rho) # inverse Euler equation
  MuVec  <- AlphaVec+ChiVec
  M <- cbind(M, matrix(MuVec,ncol = 1))             # matrix of interpolation data
  cons <- cbind(cons, matrix(ChiVec,ncol = 1))             # matrix of interpolation data
}



